package com.feidee.fd.sml.algorithm.component.feature

import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.feature.QuantileDiscretizer
import org.apache.spark.sql.DataFrame

/**
  * @author YongChen
  * @date 2019/3/22 3:18
  * @description
  * @reviewer songhaicheng
  */
case class QuantileDiscretizeEncoderParam(
                                           override val input_pt: String,
                                           override val output_pt: String,
                                           override val hive_table: String,
                                           override val flow_time: String,
                                           override val inputCol: String,
                                           override val outputCol: String,
                                           override val preserveCols: String,
                                           override val modelPath: String,
                                           // 分组数据点的桶数(分位数或类别) ，>= 2，默认值 2
                                           numBuckets: Int,
                                           // 如何处理无效值，支持 [skip, error, keep]，默认 error
                                           handleInvalid: String,
                                           // 相对误差，[0.0, 1.0]，默认值 0.001
                                           relativeError: Double
                                         ) extends FeatureParam {

  def this() = this(null, null, null, null, "input", "features", null, null , 2 , "error", 0.001)

  override def verify(): Unit = {
    super.verify()
    val handles = Array("skip", "error", "keep")
    require(handles.contains(handleInvalid.toLowerCase), s"param handleInvalid only accepts" +
      s" ${handles.mkString("[", ", ", "]")}, but has $handleInvalid")
    require(numBuckets >= 2, "param numBuckets must be greater than 1")
    require(relativeError >= 0.0 && relativeError <= 1.0, "param relativeError's range is [0.0, 1.0]")
  }
  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("numBuckets" -> numBuckets)
    map += ("handleInvalid" -> handleInvalid)
    map += ("relativeError" -> relativeError)
    map
  }
}


class QuantileDiscretizeEncoder extends AbstractFeatureEncoder[QuantileDiscretizeEncoderParam] {

  override def setUp(param: QuantileDiscretizeEncoderParam, data: DataFrame): Array[PipelineStage] = {

    val qd = new QuantileDiscretizer()
      .setInputCol(param.inputCol)
      .setOutputCol(param.outputCol)
      .setNumBuckets(param.numBuckets)
      .setHandleInvalid(param.handleInvalid)
      .setRelativeError(param.relativeError)

    Array(qd)
  }
}

object QuantileDiscretizeEncoder {

  def apply(paramStr: String): Unit = {
    new QuantileDiscretizeEncoder()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    QuantileDiscretizeEncoder(args(0))
  }

}
